# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import json
import time
import math
import tiktoken
from typing import List, Dict
from urllib.parse import urlparse
from pydantic import BaseModel, Field
from common.utils.tools import ToolsUtil
from common.utils.config import ConfigUtil
from common.chain.llm_server import BaichuanChannel
from common.chain.llm_server import OpenaiChannel
from common.chain.llm_server import BaiduChannel
from common.chain.llm_server import ZhipuChannel
from common.chain.llm_server import SparkChannel
from common.chain.llm_server import QwenChannel


class FlowsSchema(BaseModel):
    """ 流结构 """
    name: str = Field(..., max_length=100, description="场景名称")
    scene: str = Field(..., max_length=100, description="场景标识")
    alias: str = Field(..., max_length=200, description="模型别名")
    model: str = Field(..., max_length=500, description="模型名称")
    task_time: str = Field(..., max_length=100, description="任务耗时")
    use_points: int = Field(..., ge=0, description="消耗积分")
    total_tokens: int = Field(..., ge=0, description="合计Tokens")
    prompt_tokens: int = Field(..., ge=0, description="输入Tokens")
    completion_tokens: int = Field(..., ge=0, description="输出Tokens")


class ChatUtils:

    @classmethod
    def output_success(cls, idx="",
                       event: str = "finish",
                       data: str = None,
                       think: str = None,
                       model: str = "",
                       index: int = 0,
                       annex: list = None,
                       usage: dict = None,
                       stream=True,
                       refusal=None,
                       tool_calls=None,
                       function_call=None,
                       finish_reason=None,
                       **kwargs):
        """ 输出成功数据 """
        record = {
            "id": str(time.time()) if not stream and not idx else idx,
            "event": event,
            "object": "chat.completion.chunk" if stream else "chat.completion",
            "created": int(time.time()),
            "model": model,
            "choices": [
                {
                    "index": index,
                    "delta" if stream else "message": {
                        "role": "assistant",
                        "content": data,
                        "reasoning_content": think,
                        "refusal": refusal,
                        "tool_calls": tool_calls,
                        "function_call": function_call,
                        **kwargs
                    },
                    "finish_reason": finish_reason
                }
            ]
        }

        if usage:
            record["usage"] = usage
        if annex:
            record["choices"][0]["attachments"] = annex

        if not stream:
            return record

        return "data: " + json.dumps(record, ensure_ascii=False) + "\n\n"

    @classmethod
    def output_error(cls, message: str, code: any = 0, model: str = "", types=None, param=None, stream=True):
        """ 输出失败数据 """
        types = types if types else "invalid_request_error"
        record = {"error": {
            "model": model,
            "message": message,
            "code": code,
            "type": types,
            "param": param
        }}
        if not stream:
            return record
        return "data: " + json.dumps(record, ensure_ascii=False) + "\n\n"

    @classmethod
    def compute_price(cls, use_tokens: int, amount: int, points: int = 1) -> int:
        """
        计算消耗费用 (每多少tokens消耗多少积分)
        计算公式: (使用的tokens / tokens数量单位) * 收取积分数

        Args:
            use_tokens (int): 使用的tokens
            amount (int): tokens数量单位
            points (int): 收取的积分数

        Returns:
            int: 向上取整的积分数

        Author:
            zero
        """
        # 如果使用或定价为0则不消耗
        if use_tokens <= 0 or amount <= 0:
            return 0

        # 多少 amount 消耗多少 N 积分
        thousand_units = (use_tokens / amount) * points
        return math.ceil(thousand_units)

    @classmethod
    def compute_tokens(cls, text: str, model: str = "gpt-3.5-turbo"):
        try:
            encoding = tiktoken.encoding_for_model(model)
        except KeyError:
            encoding = tiktoken.get_encoding("cl100k_base")
        return len(encoding.encode(text))

    @classmethod
    def usage(cls, messages: list, reply: str = "") -> Dict[str, int]:
        reply_content: str = reply.replace("\n", "").replace(" ", "")
        completion_tokens: int = cls.compute_tokens(reply_content)

        prompt_tokens: int = 0
        for msg in messages:
            content = msg["content"] if isinstance(msg, dict) else msg[1]
            if isinstance(content, list):
                text = "".join(m["text"] for m in content if m["type"] == "text")
                text = text.replace("\n", "").replace(" ", "")
                prompt_tokens += cls.compute_tokens(text)
            else:
                text = content or ""
                text = text.replace("\n", "").replace(" ", "")
                prompt_tokens += cls.compute_tokens(text)

        return {
            "total_tokens": prompt_tokens + completion_tokens,
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens
        }


class AIChannelFactory:
    @staticmethod
    def create(channel: str, config: dict, apikey: dict = None):
        if channel == "openai":
            return OpenaiChannel(config, apikey)
        elif channel == "baichuan":
            return BaichuanChannel(config, apikey)
        elif channel == "baidu":
            return BaiduChannel(config, apikey)
        elif channel == "zhipu":
            return ZhipuChannel(config, apikey)
        elif channel == "xunfei":
            return SparkChannel(config, apikey)
        elif channel == "qwen":
            return QwenChannel(config, apikey)
        else:
            raise ValueError(f"Unsupported chat channel: {channel}")


class AIWebSearch:
    @classmethod
    async def search(cls, query: str) -> List[dict]:
        config = await ConfigUtil.get("web_search") or {}
        status: bool = bool(config.get("status") or 0)
        engine: str = str(config.get("engine") or "")
        if not status:
            return []

        if engine == "zhipu":
            return await cls.zhipu_search(query)
        else:
            raise Exception(f"Unsupported web_search channel: {engine}")

    @classmethod
    async def zhipu_search(cls, query: str) -> List[dict]:
        """
        智谱联网搜索

        Args:
            query (str): 搜索的问题

        Returns:
            List[dict]

        Author:
            zero
        """
        try:
            config = await ConfigUtil.get("web_search", "zhipu") or {}
            engine: str = config.get("search_engine", "").strip()
            apikey: str = config.get("api_key", "").strip()

            if not engine:
                raise ValueError("web_search engine is empty")
            if not apikey:
                raise ValueError("web_search apikey is empty")
            if not query:
                raise ValueError("web_search query is empty")

            results = await ToolsUtil.post_curl(
                url="https://open.bigmodel.cn/api/paas/v4/web_search",
                headers={
                    "Authorization": "Bearer " + apikey
                },
                data={
                    "search_engine": engine,
                    "search_query": query
                }
            )

            if results.get("error"):
                error = results.get("error")
                code: int = int(error.get("code"))
                message: str = str(error.get("message"))
                if code == 1703:
                    return []
                else:
                    raise Exception(f"{code}: {message} (联网搜索)")

            lists = []
            for item in results["search_result"]:
                media = item["media"]
                if not media:
                    parsed = urlparse(item["link"])
                    media = parsed.netloc

                lists.append({
                    "title": item["title"],
                    "link": item["link"],
                    "icon": item["icon"],
                    "media": media,
                    "refer": item["refer"],
                    "content": item["content"]
                })

            return lists
        except Exception as e:
            print("几个人就是就是")
            raise Exception(str(e))
